import pandas as pd
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt
import chardet
from pylab import mpl

mpl.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体：解决plot不能显示中文问题
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题


# ========= 1. 函数定义 ==========
def detect_encoding(file_path):
    """自动检测文件的编码格式"""
    with open(file_path, 'rb') as f:
        result = chardet.detect(f.read(100000))  # 读取前100KB进行检测
    return result['encoding']


def read_csv_with_encoding(file_path):
    """尝试使用多种编码格式读取CSV文件"""
    encodings = ['utf-8', 'latin1', 'cp1252']
    for enc in encodings:
        try:
            df = pd.read_csv(file_path, encoding=enc)
            print(f"成功使用编码 '{enc}' 读取文件: {file_path}")
            return df
        except UnicodeDecodeError:
            print(f"使用编码 '{enc}' 读取文件 {file_path} 失败。尝试下一个编码...")
    # 如果所有编码都失败，使用chardet自动检测
    detected_enc = detect_encoding(file_path)
    try:
        df = pd.read_csv(file_path, encoding=detected_enc)
        print(f"成功使用检测到的编码 '{detected_enc}' 读取文件: {file_path}")
        return df
    except Exception as e:
        print(f"使用检测到的编码 '{detected_enc}' 读取文件 {file_path} 失败。错误: {e}")
        raise e


# ========= 2. 数据加载 ==========
# 2.1 读取奖牌数数据
medal_counts_path = '2025_Problem_C_Data/summerOly_medal_counts.csv'
try:
    df_medals = read_csv_with_encoding(medal_counts_path)
    print("\n奖牌数数据已加载。形状:", df_medals.shape)
    # print(df_medals.head())
except Exception as e:
    print(f"无法读取奖牌数数据文件: {e}")

# 2.2 读取主办国数据
hosts_path = '2025_Problem_C_Data/summerOly_hosts.csv'
try:
    df_hosts = read_csv_with_encoding(hosts_path)
    print("\n主办国数据已加载。形状:", df_hosts.shape)
    # print(df_hosts.head())
except Exception as e:
    print(f"无法读取主办国数据文件: {e}")

# 2.3 读取赛事项目数据
programs_path = '2025_Problem_C_Data/summerOly_programs.csv'
try:
    df_programs = read_csv_with_encoding(programs_path)
    print("\n赛事项目数据已加载。形状:", df_programs.shape)
    # print(df_programs.head())
except Exception as e:
    print(f"无法读取赛事项目数据文件: {e}")

# 2.4 读取运动员数据
athletes_path = '2025_Problem_C_Data/summerOly_athletes.csv'
try:
    df_athletes = read_csv_with_encoding(athletes_path)
    print("\n运动员数据已加载。形状:", df_athletes.shape)
    # print(df_athletes.head())
except Exception as e:
    print(f"无法读取运动员数据文件: {e}")

# ========= 3. 数据清洗与整合 ==========
# 3.1 创建国家名称到NOC代码的映射字典
country_to_noc_map = {
    "United States": "USA",
    "Greece": "GRE",
    "Germany": "DEU",
    "France": "FRA",
    "Great Britain": "GBR",  # 映射为 United Kingdom
    "China": "CHN",
    "Denmark": "DNK",
    "Netherlands": "NED",
    "Hungary": "HUN",
    "Austria": "AUT",
    "Australia": "AUS",
    "Sweden": "SWE",
    "Italy": "ITA",
    "Belgium": "BEL",
    "Switzerland": "CHE",
    "Canada": "CAN",
    "Norway": "NOR",
    "Russia": "RUS",
    "Japan": "JPN",
    "Finland": "FIN",
    "Poland": "POL",
    "Turkey": "TUR",
    "Brazil": "BRA",
    "Romania": "ROU",
    "Czechoslovakia": "CZE",  # 已映射为 CZE
    "Czech Republic": "CZE",
    "South Korea": "KOR",
    "Spain": "ESP",
    "Argentina": "ARG",
    "New Zealand": "NZL",
    "South Africa": "RSA",
    "Mexico": "MEX",
    "Iran": "IRN",
    "Ethiopia": "ETH",
    "Kenya": "KEN",
    "Ukraine": "UKR",
    "Cuba": "CUB",
    "Portugal": "PRT",
    "Azerbaijan": "AZE",
    "Bahrain": "BAH",
    "Belarus": "BLR",
    "Bulgaria": "BUL",
    "Chile": "CHL",
    "Colombia": "COL",
    "Dominican Republic": "DOM",
    "Ecuador": "ECU",
    "Egypt": "EGY",
    "Estonia": "EST",
    "Georgia": "GEO",
    "Grenada": "GRD",
    "Guatemala": "GTM",
    "Hong Kong": "HKG",
    "India": "IND",
    "Indonesia": "IDN",
    "Israel": "ISR",
    "Jamaica": "JAM",
    "Jordan": "JOR",
    "Kazakhstan": "KAZ",
    "Kosovo": "KOS",
    "Kyrgyzstan": "KGZ",
    "Latvia": "LVA",
    "Lithuania": "LTU",
    "Moldova": "MDA",
    "Mongolia": "MNG",
    "Morocco": "MAR",
    "Namibia": "NAM",
    "North Korea": "PRK",
    "North Macedonia": "MKD",  # 映射为 North Macedonia
    "Norway": "NOR",
    "Pakistan": "PAK",
    "Panama": "PAN",
    "Paraguay": "PRY",
    "Peru": "PER",
    "Philippines": "PHL",
    "Qatar": "QAT",
    "Refugee Olympic Team": "ROT",
    "San Marino": "SMR",
    "Saudi Arabia": "KSA",
    "Serbia": "SRB",  # 映射为 Serbia
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Syria": "SYR",
    "Tajikistan": "TJK",
    "Thailand": "THA",
    "Trinidad and Tobago": "TTO",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Arab Emirates": "UAE",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE",
    "Antigua and Barbuda": "ATG",
    "Armenia": "ARM",
    "Bahamas": "BHS",
    "Bangladesh": "BGD",
    "Barbados": "BRB",
    "Benin": "BEN",
    "Bermuda": "BMU",
    "Bhutan": "BTN",
    "Bolivia": "BOL",
    "Bosnia and Herzegovina": "BIH",
    "Botswana": "BWA",
    "Brunei": "BRN",
    "Burkina Faso": "BFA",
    "Burundi": "BDI",
    "Cambodia": "KHM",
    "Cameroon": "CMR",
    "Cape Verde": "CPV",
    "Central African Republic": "CAF",
    "Chad": "TCD",
    "Comoros": "COM",
    "Congo (Brazzaville)": "COG",
    "Congo (Kinshasa)": "COD",
    "Costa Rica": "CRI",
    "Côte d'Ivoire": "CIV",
    "Croatia": "HRV",
    "Cyprus": "CYP",
    "Djibouti": "DJI",
    "Dominica": "DMA",
    "El Salvador": "SLV",
    "Equatorial Guinea": "GNQ",
    "Eritrea": "ERI",
    "Eswatini": "SWZ",  # 映射为 Swaziland
    "Fiji": "FJI",
    "Gabon": "GAB",
    "Gambia": "GMB",
    "Ghana": "GHA",
    "Grenada": "GRD",
    "Guinea": "GIN",
    "Guinea-Bissau": "GNB",
    "Guyana": "GUY",
    "Haiti": "HTI",
    "Honduras": "HND",
    "Iceland": "ISL",
    "Ireland": "IRL",
    "Kazakhstan": "KAZ",
    "Kenya": "KEN",
    "Kiribati": "KIR",
    "Kuwait": "KWT",
    "Kyrgyzstan": "KGZ",
    "Laos": "LAO",
    "Lebanon": "LBN",
    "Lesotho": "LSO",
    "Liberia": "LBR",
    "Libya": "LBY",
    "Liechtenstein": "LIE",
    "Luxembourg": "LUX",
    "Madagascar": "MDG",
    "Malawi": "MWI",
    "Malaysia": "MYS",
    "Maldives": "MDV",
    "Mali": "MLI",
    "Malta": "MLT",
    "Marshall Islands": "MHL",
    "Mauritania": "MRT",
    "Mauritius": "MUS",
    "Micronesia": "FSM",
    "Monaco": "MCO",
    "Montenegro": "MNE",
    "Mozambique": "MOZ",
    "Myanmar": "MMR",
    "Nauru": "NRU",
    "Nepal": "NPL",
    "Nicaragua": "NIC",
    "Niger": "NER",
    "Nigeria": "NGA",
    "Oman": "OMN",
    "Palau": "PLW",
    "Papua New Guinea": "PNG",
    "Qatar": "QAT",
    "Romania": "ROU",
    "Russia": "RUS",
    "Rwanda": "RWA",
    "Saint Kitts and Nevis": "KNA",
    "Saint Lucia": "LCA",
    "Saint Vincent and the Grenadines": "VCT",
    "Samoa": "WSM",
    "San Marino": "SMR",
    "Sao Tome and Principe": "STP",
    "Saudi Arabia": "KSA",
    "Senegal": "SEN",
    "Serbia": "SRB",  # 映射为 Serbia
    "Seychelles": "SYC",
    "Sierra Leone": "SLE",
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Solomon Islands": "SLB",
    "Somalia": "SOM",
    "South Africa": "ZAF",
    "South Korea": "KOR",
    "South Sudan": "SSD",
    "Sri Lanka": "LKA",
    "Sudan": "SDN",
    "Suriname": "SUR",
    "Sweden": "SWE",
    "Swaziland": "SWZ",  # 映射为 Eswatini
    "Switzerland": "CHE",
    "Syria": "SYR",
    "Taiwan": "TPE",  # 与noc_mapping一致
    "Tajikistan": "TJK",
    "Tanzania": "TZA",
    "Thailand": "THA",
    "Timor-Leste": "TLS",
    "Togo": "TGO",
    "Tonga": "TON",
    "Tunisia": "TUN",
    "Turkey": "TUR",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Kingdom": "GBR",
    "United States": "USA",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vanuatu": "VUT",
    "Venezuela": "VEN",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE"
}

# 3.2 清洗主办国数据，提取国家名称并映射到NOC
df_hosts['HostCountry'] = df_hosts['Host'].str.split(',').str[-1].str.strip()
df_hosts['Host_NOC'] = df_hosts['HostCountry'].map(country_to_noc_map)
df_hosts['Host_NOC'] = df_hosts['Host_NOC'].fillna('UNK')  # 'UNK' 表示未知

# 3.3 汇总运动员数据中的金牌数
# 3.3.1 过滤夏季奥运会数据
if 'Season' in df_athletes.columns:
    # 如果 'Season' 列存在，则进行过滤
    df_athletes_summer = df_athletes[df_athletes['Season'] == 'Summer'].copy()
    print("\n已基于 'Season' 列过滤夏季奥运会的运动员数据。")
else:
    # 如果 'Season' 列不存在，假设所有数据都是夏季奥运会
    df_athletes_summer = df_athletes.copy()
    print("\n假设所有运动员数据均来自夏季奥运会。")

# 3.3.2 过滤出金牌
df_athletes_gold = df_athletes_summer[df_athletes_summer['Medal'] == 'Gold'].copy()

# 3.3.3 按国家、运动项目、年份汇总金牌数
df_medal_counts = df_athletes_gold.groupby(['NOC', 'Sport', 'Year']).size().reset_index(name='Gold_Count')

# 3.4 合并主办国信息到奖牌数据
df_medal_counts = pd.merge(
    df_medal_counts,
    df_hosts[['Year', 'Host_NOC']],
    on='Year',
    how='left'
)

# 3.5 创建Host列：如果NOC == Host_NOC，则为1，否则为0
df_medal_counts['Host'] = np.where(df_medal_counts['NOC'] == df_medal_counts['Host_NOC'], 1, 0)

# 3.6 计算历史平均金牌数
df_historical = df_medal_counts.groupby('NOC')['Gold_Count'].mean().reset_index().rename(
    columns={'Gold_Count': 'HistoricalGold'})

# 3.7 合并历史金牌数到奖牌数据
df_final = pd.merge(df_medal_counts, df_historical, on='NOC', how='left')

# 3.8 处理特殊NOC
# 移除 'mixed team' 和其他特殊NOC
df_final = df_final[~df_final['NOC'].isin(['MIX'])]

# 3.9 计算每年的总赛事数
# 转换赛事项目数据为长格式
df_programs_melt = df_programs.melt(
    id_vars=['Sport', 'Discipline', 'Code', 'Sports Governing Body'],
    var_name='Year',
    value_name='Num_Events'
)

# 提取年份数字，并去除可能的非数字字符
df_programs_melt['Year'] = pd.to_numeric(df_programs_melt['Year'], errors='coerce')

# 处理 'Num_Events' 列，转换为数值型，非数值字符如 '•' 转换为 0
df_programs_melt['Num_Events'] = pd.to_numeric(df_programs_melt['Num_Events'], errors='coerce').fillna(0)

# 计算每年的Total_Events
df_total_events = df_programs_melt.groupby('Year')['Num_Events'].sum().reset_index().rename(
    columns={'Num_Events': 'Total_Events'})

# 3.10 合并Total_Events到df_final
df_final = pd.merge(df_final, df_total_events, on='Year', how='left')

# 填补缺失的Total_Events为0
df_final['Total_Events'] = df_final['Total_Events'].fillna(0)

# ========= 4. 定义“伟大教练”名单并标注 ==========
# 4.1 定义“伟大教练”名单
# 定义“伟大教练”的执教记录
great_coaches = [
    {
        'Coach_Name': 'Lang Ping',
        'NOC': 'CHN',
        'Sport': 'Volleyball',
        'Year': 2021  # 东京奥运会
    },
    {
        'Coach_Name': 'Lang Ping',
        'NOC': 'CHN',
        'Sport': 'Volleyball',
        'Year': 2016  # 里约奥运会
    },
    {
        'Coach_Name': 'Lang Ping',
        'NOC': 'USA',
        'Sport': 'Volleyball',
        'Year': 2008  # 北京奥运会
    },
    {
        'Coach_Name': 'Bola Karoly',
        'NOC': 'USA',
        'Sport': 'Gymnastics',
        'Year': 1996  # 亚特兰大奥运会
    },
    {
        'Coach_Name': 'Bola Karoly',
        'NOC': 'ROU',
        'Sport': 'Gymnastics',
        'Year': 1976  # 蒙特利尔奥运会
    },
    {
        'Coach_Name': 'Pete Hart',
        'NOC': 'USA',
        'Sport': 'Swimming',
        'Year': 1988  # 首尔奥运会
    },
    {
        'Coach_Name': 'Pete Hart',
        'NOC': 'AUS',
        'Sport': 'Swimming',
        'Year': 2000  # 悉尼奥运会
    },
    {
        'Coach_Name': 'Jos Gonzalez',
        'NOC': 'ESP',
        'Sport': 'Wrestling',
        'Year': 1988  # 首尔奥运会
    },
    {
        'Coach_Name': 'Jos Gonzalez',
        'NOC': 'ARG',
        'Sport': 'Fencing',
        'Year': 2004  # 雅典奥运会
    },
    {
        'Coach_Name': 'Igor Muleshoe',
        'NOC': 'RUS',
        'Sport': 'Gymnastics',
        'Year': 2000  # 悉尼奥运会
    },
    {
        'Coach_Name': 'Igor Muleshoe',
        'NOC': 'BLR',
        'Sport': 'Gymnastics',
        'Year': 2004  # 雅典奥运会
    },
    {
        'Coach_Name': 'Howard Ferrell',
        'NOC': 'USA',
        'Sport': 'Basketball',
        'Year': 1972  # 慕尼黑奥运会
    },
    {
        'Coach_Name': 'Howard Ferrell',
        'NOC': 'AUS',
        'Sport': 'Basketball',
        'Year': 2000  # 悉尼奥运会
    },
    {
        'Coach_Name': 'Howard Ferrell',
        'NOC': 'NZL',
        'Sport': 'Basketball',
        'Year': 2008  # 北京奥运会
    },
    {
        'Coach_Name': 'Dumont Jackson',
        'NOC': 'USA',
        'Sport': 'Athletics',
        'Year': 1968  # 墨西哥城奥运会
    },
    {
        'Coach_Name': 'Dumont Jackson',
        'NOC': 'KEN',
        'Sport': 'Athletics',
        'Year': 1980  # 莫斯科奥运会
    },
    {
        'Coach_Name': 'Neil Adams',
        'NOC': 'GBR',
        'Sport': 'Judo',
        'Year': 1980  # 莫斯科奥运会
    },
    {
        'Coach_Name': 'Neil Adams',
        'NOC': 'USA',
        'Sport': 'Judo',
        'Year': 1984  # 洛杉矶奥运会
    },
    {
        'Coach_Name': 'Werner Grave',
        'NOC': 'GER',
        'Sport': 'Football',
        'Year': 1972  # 慕尼黑奥运会
    },
    {
        'Coach_Name': 'Werner Grave',
        'NOC': 'GHA',
        'Sport': 'Football',
        'Year': 1982  # 洛杉矶奥运会
    },
    {
        'Coach_Name': 'Glenn Hopkins',
        'NOC': 'USA',
        'Sport': 'Volleyball',
        'Year': 1984  # 洛杉矶奥运会
    },
    {
        'Coach_Name': 'Glenn Hopkins',
        'NOC': 'BRA',
        'Sport': 'Volleyball',
        'Year': 2008  # 北京奥运会
    },
    {
        'Coach_Name': 'Franz Rickenbacker',
        'NOC': 'FRG',
        'Sport': 'Football',
        'Year': 1974  # 西德世界杯（非奥运会）
    },
    {
        'Coach_Name': 'Franz Rickenbacker',
        'NOC': 'ARG',
        'Sport': 'Football',
        'Year': 1968  # 墨西哥奥运会
    },
    {
        'Coach_Name': 'Alfredo Di Stefano',
        'NOC': 'ESP',
        'Sport': 'Football',
        'Year': 1968  # 墨西哥奥运会
    },
    {
        'Coach_Name': 'Alfredo Di Stefano',
        'NOC': 'COL',
        'Sport': 'Football',
        'Year': 1992  # 巴塞罗那奥运会
    },
    {
        'Coach_Name': 'Mirko Novell',
        'NOC': 'ITA',
        'Sport': 'Athletics',
        'Year': 1992  # 巴塞罗那奥运会
    },
    {
        'Coach_Name': 'Mirko Novell',
        'NOC': 'NLD',
        'Sport': 'Athletics',
        'Year': 2000  # 悉尼奥运会
    },
    {
        'Coach_Name': 'Thomas Bach',
        'NOC': 'GER',
        'Sport': 'Fencing',
        'Year': 1988  # 汉城奥运会
    },
    {
        'Coach_Name': 'Thomas Bach',
        'NOC': 'FRA',
        'Sport': 'Fencing',
        'Year': 1996  # 亚特兰大奥运会
    },
    {
        'Coach_Name': 'Kevin Keegan',
        'NOC': 'GBR',
        'Sport': 'Football',
        'Year': 1992  # 巴塞罗那奥运会
    },
    {
        'Coach_Name': 'Kevin Keegan',
        'NOC': 'AUS',
        'Sport': 'Football',
        'Year': 2000  # 悉尼奥运会
    },
    {
        'Coach_Name': 'Aldrich Ola',
        'NOC': 'RSA',
        'Sport': 'Basketball',
        'Year': 1996  # 亚特兰大奥运会
    },
    {
        'Coach_Name': 'Aldrich Ola',
        'NOC': 'MAR',
        'Sport': 'Basketball',
        'Year': 2004  # 雅典奥运会
    },
    {
        'Coach_Name': 'George Cavil',
        'NOC': 'GBR',
        'Sport': 'Swimming',
        'Year': 1956  # 墨尔本奥运会
    },
    {
        'Coach_Name': 'George Cavil',
        'NOC': 'BRA',
        'Sport': 'Swimming',
        'Year': 1992  # 巴塞罗那奥运会
    },
    {
        'Coach_Name': 'Jack Rahman',
        'NOC': 'USA',
        'Sport': 'Diving',
        'Year': 1984  # 洛杉矶奥运会
    },
    {
        'Coach_Name': 'Jack Rahman',
        'NOC': 'FRA',
        'Sport': 'Diving',
        'Year': 1996  # 亚特兰大奥运会
    },
    {
        'Coach_Name': 'Carl Strauss',
        'NOC': 'GER',
        'Sport': 'Volleyball',
        'Year': 1972  # 慕尼黑奥运会
    },
    {
        'Coach_Name': 'Carl Strauss',
        'NOC': 'TCH',
        'Sport': 'Volleyball',
        'Year': 1980  # 莫斯科奥运会
    },
    {
        'Coach_Name': 'Robert Coleman',
        'NOC': 'CAN',
        'Sport': 'Ice Hockey',
        'Year': 1976  # 蒙特利尔奥运会
    },
    {
        'Coach_Name': 'Robert Coleman',
        'NOC': 'AUS',
        'Sport': 'Ice Hockey',
        'Year': 2002  # 盐湖城冬季奥运会
    }
    # 添加更多伟大教练的信息
]

df_coaches = pd.DataFrame(great_coaches)
print("\n伟大教练名单:")
print(df_coaches)

# 4.2 标注“伟大教练”效应
# 先初始化Coach列为0
df_final['Coach'] = 0

# 遍历每个“伟大教练”的记录，并标注相应的记录
for index, coach in df_coaches.iterrows():
    condition = (
            (df_final['NOC'] == coach['NOC']) &
            (df_final['Sport'] == coach['Sport']) &
            (df_final['Year'] == coach['Year'])
    )
    df_final.loc[condition, 'Coach'] = 1

# 设置显示所有列
pd.set_option('display.max_columns', None)
print("\n标注“伟大教练”后的奖牌数据:")

# print(df_final.head(15))

# ========= 5. 构建并拟合负二项回归模型 ==========
# 5.1 添加运动项目和年份的固定效应
df_final['Sport_FE'] = df_final['Sport']
df_final['Year_FE'] = df_final['Year']  # 确保 Year_FE 是数值型

# 检查 Year_FE 的数据类型
print("\nYear_FE 的数据类型:")
print(df_final[['Year_FE']].dtypes)

# 5.2 构建模型公式
# 为了让模型能够泛化到新NOCs，移除 Country_FE 作为固定效应
# 将 Year_FE 作为数值型变量处理，而非分类变量
formula = 'Gold_Count ~ Coach + C(Sport_FE) + Year_FE'

# 5.3 拟合泊松回归模型
model_poisson = smf.glm(formula=formula, data=df_final, family=sm.families.Poisson()).fit()
print("\n泊松回归模型结果:")
print(model_poisson.summary())

# 5.4 检查过度分散
mean_count = df_final['Gold_Count'].mean()
var_count = df_final['Gold_Count'].var()
dispersion = var_count / mean_count
print(f"\n均值 (Mean Count): {mean_count}")
print(f"方差 (Variance Count): {var_count}")
print(f"离散比率 (Dispersion Ratio - Variance/Mean): {dispersion}")

if dispersion > 1.5:
    print("数据存在过度分散，建议使用负二项回归模型。")
    # 5.5 拟合负二项回归模型
    model_nb = smf.glm(formula=formula, data=df_final, family=sm.families.NegativeBinomial()).fit()
    print("\n负二项回归模型结果:")
    print(model_nb.summary())
else:
    model_nb = model_poisson

# ========= 6. 模型性能评估 ==========
# 6.1 计算模型在训练数据上的预测
df_final['Predicted'] = model_nb.predict(df_final)

# 6.2 计算性能指标
mse = mean_squared_error(df_final['Gold_Count'], df_final['Predicted'])
mae = mean_absolute_error(df_final['Gold_Count'], df_final['Predicted'])
r2 = r2_score(df_final['Gold_Count'], df_final['Predicted'])

print(f"\n模型性能评估：\n均方误差 (MSE)：{mse:.2f}\n平均绝对误差 (MAE)：{mae:.2f}\n决定系数 (R²)：{r2:.2f}")

# 生成预测数据，分别为有和没有“伟大教练”的情况
years = np.arange(1996, 2025)
sports = 'Volleyball'  # 排球
noc_china = 'CHN'

# 获取Total_Events的均值或历史数据
average_total_events = df_final['Total_Events'].mean()

# 预测没有伟大教练（Coach=0）和有伟大教练（Coach=1）时的金牌数
predictions_no_coach = []
predictions_with_coach = []

for year in years:
    # 预测没有伟大教练时的金牌数
    predicted_gold_no_coach = model_nb.predict(pd.DataFrame({
        'NOC': [noc_china],
        'Sport_FE': [sports],
        'Coach': [0],
        'Year_FE': [year],
        'Total_Events': [average_total_events]
    }))[0]

    # 预测有伟大教练时的金牌数
    predicted_gold_with_coach = model_nb.predict(pd.DataFrame({
        'NOC': [noc_china],
        'Sport_FE': [sports],
        'Coach': [1],
        'Year_FE': [year],
        'Total_Events': [average_total_events]
    }))[0]

    predictions_no_coach.append(predicted_gold_no_coach)
    predictions_with_coach.append(predicted_gold_with_coach)

# 绘制图表
plt.figure(figsize=(10, 6))
plt.plot(years, predictions_no_coach, label='No Great Coach', marker='o', color='orange', linestyle='--')
plt.plot(years, predictions_with_coach, label='There are great coaches', marker='o', color='blue', linestyle='-')

# 添加标签和标题
plt.xlabel('Year', fontsize=12)
plt.ylabel('Gold Count', fontsize=12)
plt.title('“The impact of “great coaches” on the number of gold medals (Volleyball, China)', fontsize=14)
plt.xticks(years, rotation=45)
plt.grid(True, linestyle='--', alpha=0.5)
plt.legend()

plt.tight_layout()
plt.show()

